package com.smh.aiagent.config;

import org.springframework.ai.model.tool.DefaultToolCallingManager;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.execution.DefaultToolExecutionExceptionProcessor;
import org.springframework.ai.tool.execution.ToolExecutionException;
import org.springframework.ai.tool.execution.ToolExecutionExceptionProcessor;
import org.springframework.ai.tool.resolution.DelegatingToolCallbackResolver;
import org.springframework.ai.tool.resolution.StaticToolCallbackResolver;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;

import java.io.IOException;
import java.util.List;

/**
 * @author shiminghui
 * @date 2025/6/5 16:11
 * @description: TODO
 */
@Configuration
public class ToolCallingManagerConfig {


    /**
     * 工具调用管理器
     *
     * @param toolExecutionExceptionProcessor
     * @return
     */
    @Bean
    public ToolCallingManager toolCallingManager(ToolExecutionExceptionProcessor toolExecutionExceptionProcessor, ToolCallback[] allTools) {
        return DefaultToolCallingManager
                .builder()
//                .observationRegistry() //  设置观察者注册器
                .toolCallbackResolver(new DelegatingToolCallbackResolver(List.of(
                        new StaticToolCallbackResolver(List.of(allTools)),
                        new StaticToolCallbackResolver(List.of(allTools))
                ))) // 设置工具回调解析器
                .toolExecutionExceptionProcessor(toolExecutionExceptionProcessor) // 设置工具执行异常处理器
                .build();
    }

    /**
     * 工具执行异常处理器
     *
     * @return
     */
    @Bean
    @Primary
    public ToolExecutionExceptionProcessor toolExecutionExceptionProcessor() {
        return DefaultToolExecutionExceptionProcessor
                .builder()
                .alwaysThrow(true) // 设置是否总是抛出异常,(抛出异常就是直接退出,不抛出就是打个日志,然后将错误信息发给AI)
//                .wait(1000)
                .build();
    }

    @Bean
    public ToolExecutionExceptionProcessor customToolExecutionExceptionProcessor() {
        return new ToolExecutionExceptionProcessor() {
            @Override
            public String process(ToolExecutionException exception) {
                if (exception.getCause() instanceof IOException) {
                    return "Unable to access external resource. Please try a different approach.";
                } else if (exception.getCause() instanceof SecurityException) {
                    throw exception;
                }
                return "Error executing tool: " + exception.getMessage();
            }
        };
    }
}
